Mocking around with Python and Unittest

As a System Engineer I have encountered multiple different APIs throughout my career. Some of them I created myself, some of them I consumed in my code. And because I'm a strong advocate for writing tests I had to mock them. Code in this article will use Python's unittest module and its mock module.

The MI6 Secret Agent API

We will be testing below function. The function calls MI6's API to get the list of its agents. Try not to tell about it to anyone, it's a secret API.

# my_module
import requests

def call_api(data):
    try:
        response = requests.post("http://mi6.uk/agents/", data=data)
        if response.status_code == 200:
            print("Success")
        return response.json()
    except requests.exception.RequestException as e:
        print("Error occurred: {}".format(e))
        return {}

Tests that call external APIs should execute regardless if the environment they are run on is connected to the internet. If we want to test the functionalities that depend on the API call results we should make sure that our tests will not rely on the internet availability. Shortly speaking, tests should pass everywhere.

Enter Mock module

With mock module we can specify a value that above requests.post method will return. Here we have a test where we call my_module.call_api function, requesting the list of secret agents. To keep things short, it will return only data of single agent.

If you look at the call_api function, you have couple ways to mock calling the API. We could mock entire requests module or just requests.post method or even just json method on the object returned by the post method.

Because aside from manipulating the API result I would like control status code on the HTTP response I decided to mock the post method. The requests.post is what returns the HTTP response object, that will allow us to easily manipulate object mocking the response.

# tests.py
from unittest import TestCase, main, mock

from my_module import call_api

class MyModuleTest(TestCase):
    @mock.patch("my_module.requests.post")                #1
    def test_call_api(self, mock_post):                   #2
        my_mock_response = mock.Mock(status_code=200)     #3
        my_mock_response.json.return_value = {            #4
            "result": [
                {
                    "name": "James Bond",
                    "is_on_mission": True
                }
            ]
        }
        mock_post.return_value = my_mock_response         #5

        response = call_api(data={"secret": True})
        self.assertEqual(response.status_code, 200)

        agent_data = response["result"][0]
        self.assertEqual(agent_data["name"], "James Bond")

if __name__ == "__main__":
    main()

There are 5 important points I want you to pay attention to.

#1 - Using the mock.patch as a decorator we specify the target of our mocking - requests.post. This is the most important thing to understand in using mock's patch decorator. In order to force the 'post' method return our value it has to be mocked in the module that is using the 'post' method. In our example the requests.post is called in my_module so the path "my_module.requests.post" must be the first argument to the patch decorator. That is the most confusing part for beginners. The mocking happens in the module where target function or object is used. If we provided patch decorator with just requests.post, we would not get the expected value in our code.

#2 When you use patch method as a decorator, mocks are passed to the decorated function which allows you to do some more magic with them, I'll explain that later in the article. For now you need to know that your unittest function requires an additional argument, I called it mock_post, this way I know right away what kind of object it is.

#3 We will make use of Mock class to imitate the HTTP response. requests.post method returns object that has status_code attribute and json method. The way we instantiate the Mock class here, we set the status_code attribute to be equal to 200, so when the code in the call_api function will access it, the attribute will return integer with value equal 200.

#4 This is where you assign a value that mocked post method returns. We know from not existing specification that the API on successful response returns json with "result" key, and a list of jsons, each of which contains data of single agent. As mentioned in #3, response object has the json method that returns Python's dictionary. We can attach on the fly any number of methods and attributes to the Mock instance as we like. Again, to imitate the HTTP response's json method we just access the my_mock_response with dot notation and method name. Here unittest.mock creates new Mock object on the my_mock_response.json attribute. This way if the code in call_api function will call json on the response object, it will get a the mocked API result.

#5 Lastly here we set the value that we want the post method to return - the mocked HTTP response. Without this line, mock would return instance of unittest.mock.Mock class.

If you don't set anything, the post method will return new Mock object on the fly. That is the main behavior of this module. You could access any attribute on that object, even the one that doesn't exist and it will be added to the Mock object.

If you run the above test, it will finish with single successful test. We verified that the response's status code is 200, and acccessed the only result containing "James Bond" name. In my_module.py the print function would print "Success" and call_api return the response's json as a dictionary.

There is a little problem with call_api function. Regardless of the status_code on the response, it calls json method on it. That means if the endpoint has changed we might end up with 404 response and in that case an exception would be thrown inside call_api because json method was called on the HTTP response that might not have any json data in the first place. This is dangerous gotcha that comes with mocking. The mock module gives you great flexibility, but leaves you with all the burden of imitating the exact behavior of system you are working on.

Now that we know we might not get json in the response (perhaps our boss is yelling that production code has the bug - never happened to me) we want to write a test case with that exact behavior to prevent from happening again.

The updated, non-existing specification said that the call_api function should return the exception's message that comes with the exception raised when calling json method.

# my_module
import requests

def call_api(data):
    try:
        response = requests.post("http://mi6.uk/agents/", data=data)
        try:
            return response.json()
        except Exception as e:
            return e.args[0]
    except requests.exception.RequestException as e:
        print("Error occurred: {}".format(e))
        return {}

How can we cause an exception to be raised in our tests?

The side_effect attribute

From documentation we read [side_effect] ~ Useful for raising exceptions or dynamically changing return values.

Let's try it in test.

    @mock.patch("my_module.requests.post")
    def test_call_api_no_json(self, mock_post):
        my_mock_response = mock.Mock(status_code=404)
        my_mock_response.json.side_effect = ValueError("here be dragons")     #1
        mock_post.return_value = my_mock_response

        response = call_api(data={"secret": True})
        self.assertEqual(response, "here be dragons")

#1 If you compare this test with previous one, the main difference happens on the mocked json method. Here instead of using return_value, we use special side_effect attribute and assign ValueError with message of our choice. Next time the code in my_module calls json method that belongs to the object returned by requests.post method, ValueError will be raised. If we run the tests now, we get 2 successful test cases veryfing that response is indeed what we expect.

side_effect has one more very interesting use case. We can indeed dynamically change return values and exceptions interchangeably. Let's say the specification has changed again. Instead of returning the error message, keep calling the API until you get the result.

import requests

def call_api(data):
    try:
    	while True:
	        response = requests.post("http://mi6.uk/agents/", data=data)
	        try:
	        	return response.json()
	        except Exception as e:
	        	print(e.args[0])
	        	continue
    except Exception as e:
        print("Error occurred: {}".format(e))
        return {}

Let's assume this scenario: - first two calls to the API will return responses without json - the third one will return a response that has a json

@mock.patch("my_module.requests.post")
def test_call_api_until_there_is_a_result(self, mock_post):
    result = {
        "result": [
            {
                "name": "James Bond",
                "is_on_mission": True
            }
        ]
    }

    my_mock_response = mock.Mock(status_code=200)
    my_mock_response.json.side_effect = [
        ValueError("these are not the results you are looking for"),
        ValueError("use force..."),
        result
    ]
    mock_post.return_value = my_mock_response

    response = call_api(data={"secret": True})
    agent_data = response["result"][0]
    self.assertEqual(agent_data["name"], "James Bond")

To keep the above test code short, instead of preparing separate response for each call, I'm using the side_effect attribute on json method of the response. If you wanted to create different responses with different status codes and so on, you could use side_effect in the same way buy on the mocked requests.post method, like that: my_mock_response.side_effect = [...]. In the iterator you provide multiple Mock(status_code=400, ...)

To make our scenario happen, we assign an iterator to the side_effect attribute. Each time the json method will be called, a value from the list above will be returned. First call results in ValueError and "these are not the results you are looking for" is printed out. Second call again, no luck, "use force..." is printed, let's try again. Third time the charm, we've got it! The data is returned and we can verify the name of our agent is James Bond.

Thing to remember: if we intended to call json method on that response 4th time we would end up with StopIteration error. We will keep getting results until there are any values left with provided iterator, and because we had only 3 values, the 4th call results in error. We can clear side_effect by setting it to None which would make the json method start returning Mock objects.

Mock's methods and attributes

According to our previous test case, the ValueError should be raised twice. We can verify that using the object passed to the test function with mock.patch decorator I mentioned earlier. If you look at our call_api function, in case an exception being raised while calling json method, its message is printed out. Until now we could have seen it in the console output.

To do that we can mock the print function called in the except... block. We add second mock.patch above the decorator that is mocking requests.post method and another argument to our test function. The order of arguments and decorators matters. The first decorator applied to test function inserts first argument, second decorator inserts second argument, and so on. Imagine decorators walking into the function from above, turning right and leaving an argument on the way out lined up.

@mock.patch("my_module.print")
@mock.patch("my_module.requests.post")
def test_call_api_until_there_is_a_result(self, mock_post, mock_print):

This allows as to access interesting methods and attributes that Mock object has. Some of them start with assert prefix. For the whole list check the documentation. You can verify that the mocked function has been called once, multiple times, the arguments it has been called with or if it was called at all.

For example if we call the mock_print.assert_called() in our test, nothing happens in particular. The reason behind this is that the print function has been called in the call_api and the verification passes. But if we call mock_print.assert_called_once() we get the errror: AssertionError: Expected 'print' to have been called once. Called 2 times.

We can also verify what arguments that the mocked function has been called with by accessing call_args or call_args_list. The difference between those is the first has a value of the last call of the mocked function and the latter has values of all calls. But they don't return plain values immediately. They return the call object that is wrapping the arguments.

# tests.py
mock_print.assert_called()   # OK
mock_print.assert_called_once()   # AssertionError

# Test first call of print function
self.assertEqual(mock_print.call_args_list[0][0][0], "these are not the results you are looking for")
# Test second call of print function
self.assertEqual(mock_print.call_args_list[1][0][0], "use force...")
mock_print.assert_called_with(("use force..."))

If you look back to the console, you can notice that the above values are printed into the console anymore. That should be obvious because we are mocking the print function which is now an instance of Mock class in my_module.py

To mock function but keep original functionality we must pass replacement to the mock.patch decorator with 'wraps' argument - the print function.

@mock.patch("my_module.print", wraps=print)
@mock.patch("my_module.requests.post")
def test_call_api_until_there_is_a_result(self, mock_post, mock_print):

Run the tests again, and notice it works again. We have verified how many times the print was called, what arguments have been used for the calls.

Summary

We have explored some basic features that mock module provides and learned basic mocking strategy, where to mock, what to mock and how verify results of the operations. The module allows for vast range of options when it comes to testing our code but we must be aware of the behavior that we try to imitate. Mock allows us to write tests that can be run on any environment even one without internet access which is especially valuable if your code depends on external APIs.